// main.go（修复版）
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"time"
)

type Message struct {
	Cmd  string `json:"cmd"`
	Data string `json:"data"`
}

func main() {
	const (
		totalConnections = 100
		msgInterval      = 100 * time.Millisecond
		testDuration     = 30 * time.Second
	)

	var wg sync.WaitGroup
	var success, failed int64
	// var mu sync.Mutex

	// 使用 context 控制超时
	ctx, cancel := context.WithTimeout(context.Background(), testDuration)
	defer cancel()

	// 监听 Ctrl+C
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		fmt.Println("\n🛑 中断信号，正在关闭...")
		cancel()
	}()

	for i := 0; i < totalConnections; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()

			conn, err := net.Dial("tcp", "localhost:8080")
			if err != nil {
				log.Printf("连接失败: %v", err)
				atomicAdd(&failed, 1)
				return
			}
			defer conn.Close()

			ticker := time.NewTicker(msgInterval)
			defer ticker.Stop()

			for {
				select {
				case <-ctx.Done():
					return // 超时或中断

				case <-ticker.C:
					// 设置写超时，避免阻塞
					_ = conn.SetWriteDeadline(time.Now().Add(1 * time.Second))

					msg := Message{
						Cmd:  "ping",
						Data: fmt.Sprintf("client-%d", id),
					}
					data, _ := json.Marshal(msg)
					_, err := conn.Write(append(data, '\n'))
					if err != nil {
						atomicAdd(&failed, 1)
						return // 写失败，退出
					}
					atomicAdd(&success, 1)
				}
			}
		}(i)
	}

	fmt.Printf("🚀 开始压测，持续 %v...\n", testDuration)
	<-ctx.Done() // 等待超时或中断

	fmt.Println("⏱️  压测结束，等待连接关闭...")
	wg.Wait() // 等待所有 goroutine 退出

	// 取消可能还在运行的 ticker
	cancel()

	fmt.Printf("✅ 成功发送: %d 次\n", load(&success))
	fmt.Printf("❌ 失败: %d 次\n", load(&failed))
}

// 线程安全的原子操作
func atomicAdd(ptr *int64, val int64) {
	atomic.AddInt64(ptr, val)
}

func load(ptr *int64) int64 {
	return atomic.LoadInt64(ptr)
}
